# 模型微调
from keras.applications import VGG16
from keras import models, optimizers
from keras import layers
from keras.preprocessing.image import ImageDataGenerator
import os
import matplotlib.pyplot as plt


'''
weights='imagenet',  模型初始化的权重检查点
include_top=False,   指定模型最后是否包含密集连接分类器。
默认情况下，这个密集连接分类器对应于ImageNet的1000个类别。这里我们使用的就是两个类别（cats and dogs）
input_shape=(150, 150, 3)  输入到网络中的图像张量形状。
'''
# path = "../mydatas/vgg16_weights_tf_dim_ordering_tf_kernels_notop (1).h5"

conv_base = VGG16(weights='imagenet',
                  include_top=False,
                  input_shape=(150, 150, 3))

model = models.Sequential()
model.add(conv_base)
model.add(layers.Flatten())
model.add(layers.Dense(256, activation='relu'))
model.add(layers.Dense(1, activation='sigmoid'))
print(conv_base.summary())


'''
调整最后三个卷积层，就是直到block4_pool的所有层都应该被冻结，
                而block5_conv1、block5_conv2和block_conv3这三层是可训练的
'''

# 冻结直到某一层的所有层
conv_base.trainable = True
# 只保留卷积基最后三层为解冻状态，其他层重新锁定
for layer in conv_base.layers[:15]:
    layer.trainable = False


# 微调模型
base_dir = "G:\\python\\keshe\\train1"
# base_dir = 1
print(base_dir)
train_dir = os.path.join(base_dir, 'train')
validation_dir = os.path.join(base_dir, 'validation')
test_dir = os.path.join(base_dir, 'test')


train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,# 角度值(0-180)表示图像随机旋转的范围
    width_shift_range=0.2,# 水平向上平移的范围(相对于总宽度的比例)
    height_shift_range=0.2, # 垂直向上平移的范围(相对于总高度的比例)
    shear_range=0.2, # 随机错切变换的角度
    zoom_range=0.2, # 随机缩放的范围
    horizontal_flip=True, # 随机将图像水平翻转
    fill_mode='nearest'# 用于填充新创建像素的方法，可能来自于旋转或宽度/高度平移

)

test_datagen = ImageDataGenerator(rescale=1./255)

train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(150, 150),
    batch_size=40,
    class_mode='binary'
)

validation_generator = test_datagen.flow_from_directory(
    validation_dir,
    target_size=(150, 150),
    batch_size=40,
    class_mode='binary'
)


model.compile(loss='binary_crossentropy',
                 optimizer=optimizers.RMSprop(lr=1e-5),  # 之所以让学习率很小，是因为对于微调的三层表示，我们希望变化范围不要太大。
              # 太大的权重更新可能会破坏这些。
                 # optimizer='rmsprop',
                 metrics=['acc'])

h = model.fit_generator(
    train_generator,
    steps_per_epoch=200,
    epochs=100,
    validation_data=validation_generator,
    validation_steps=50
)

model.save('./data/cats_and_dogs_wei_tiao.h5')
# model.save('./data/cats_and_dogs2.h52')

# 绘制损失率和准确率图像
acc = h.history.get('acc')
val_acc = h.history.get('val_acc')
loss = h.history.get('loss')
val_loss = h.history.get('val_loss')

epochs = range(1, len(acc) + 1)

plt.plot(epochs, acc, 'r', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.xlabel('epochs')
plt.ylabel('per')
plt.legend()

plt.figure()

plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.xlabel('epochs')
plt.ylabel('per')
plt.legend()

plt.show()


